package com.emc.greenplum.pigsterix.compiler;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import org.apache.commons.lang3.mutable.Mutable;
import org.apache.commons.lang3.mutable.MutableObject;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.pig.impl.PigContext;
import org.apache.pig.impl.logicalLayer.FrontendException;
import org.apache.pig.newplan.DependencyOrderWalker;
import org.apache.pig.newplan.Operator;
import org.apache.pig.newplan.OperatorPlan;
import org.apache.pig.newplan.logical.expression.LogicalExpressionPlan;
import org.apache.pig.newplan.logical.expression.ProjectExpression;
import org.apache.pig.newplan.logical.relational.LOCogroup;
import org.apache.pig.newplan.logical.relational.LOCross;
import org.apache.pig.newplan.logical.relational.LODistinct;
import org.apache.pig.newplan.logical.relational.LOFilter;
import org.apache.pig.newplan.logical.relational.LOForEach;
import org.apache.pig.newplan.logical.relational.LOGenerate;
import org.apache.pig.newplan.logical.relational.LOInnerLoad;
import org.apache.pig.newplan.logical.relational.LOJoin;
import org.apache.pig.newplan.logical.relational.LOLimit;
import org.apache.pig.newplan.logical.relational.LOLoad;
import org.apache.pig.newplan.logical.relational.LONative;
import org.apache.pig.newplan.logical.relational.LOSort;
import org.apache.pig.newplan.logical.relational.LOSplit;
import org.apache.pig.newplan.logical.relational.LOSplitOutput;
import org.apache.pig.newplan.logical.relational.LOStore;
import org.apache.pig.newplan.logical.relational.LOStream;
import org.apache.pig.newplan.logical.relational.LOUnion;
import org.apache.pig.newplan.logical.relational.LogicalRelationalNodesVisitor;
import org.apache.pig.newplan.logical.relational.LogicalRelationalOperator;
import org.apache.pig.newplan.logical.relational.LogicalSchema;
import org.apache.pig.newplan.logical.relational.LogicalSchema.LogicalFieldSchema;

import edu.uci.ics.hyracks.algebricks.core.algebra.base.ILogicalExpression;
import edu.uci.ics.hyracks.algebricks.core.algebra.base.ILogicalOperator;
import edu.uci.ics.hyracks.algebricks.core.algebra.base.ILogicalPlan;
import edu.uci.ics.hyracks.algebricks.core.algebra.base.LogicalVariable;
import edu.uci.ics.hyracks.algebricks.core.algebra.expressions.ScalarFunctionCallExpression;
import edu.uci.ics.hyracks.algebricks.core.algebra.expressions.VariableReferenceExpression;
import edu.uci.ics.hyracks.algebricks.core.algebra.functions.AlgebricksBuiltinFunctions;
import edu.uci.ics.hyracks.algebricks.core.algebra.metadata.IDataSink;
import edu.uci.ics.hyracks.algebricks.core.algebra.operators.logical.DataSourceScanOperator;
import edu.uci.ics.hyracks.algebricks.core.algebra.operators.logical.EmptyTupleSourceOperator;
import edu.uci.ics.hyracks.algebricks.core.algebra.operators.logical.InnerJoinOperator;
import edu.uci.ics.hyracks.algebricks.core.algebra.operators.logical.WriteOperator;
import edu.uci.ics.hyracks.algebricks.core.algebra.plan.ALogicalPlanImpl;

public class Pig2AlgebricksVisitor extends LogicalRelationalNodesVisitor {

	public static int currentVariable = 0;

	/**
	 * map variable name to the logical variable
	 */
	private HashMap<String, LogicalVariable> nameToLogicalVariableMap = new HashMap<String, LogicalVariable>();
	/**
	 * map field name to LogicalVariable
	 */
	private HashMap<String, LogicalVariable> fieldToLogicalVariableMap = new HashMap<String, LogicalVariable>();

	/**
	 * map logical variable to name
	 */
	private HashMap<LogicalVariable, String> logicalVariableToFieldMap = new HashMap<LogicalVariable, String>();

	private List<Mutable<ILogicalOperator>> roots;

	// String = alias (A), maps to operatorTree (ILogicalOperator) + schema (map
	// structure)
	private Map<String, Pair<ILogicalOperator, Map<String, LogicalVariable>>> currentPlan = new HashMap<String, Pair<ILogicalOperator, Map<String, LogicalVariable>>>();
	protected PigContext pc;

	public Pig2AlgebricksVisitor(OperatorPlan plan) throws FrontendException {
		super(plan, new DependencyOrderWalker(plan));
		roots = new ArrayList<Mutable<ILogicalOperator>>();
	}

	public void setPigContext(PigContext pc) {
		this.pc = pc;
	}

	public void visit(LOLoad load) throws FrontendException {

		LogicalSchema schema = load.getSchema();
		List<LogicalFieldSchema> fields = schema.getFields();

		List<Byte> types = new ArrayList<Byte>();

		ArrayList<LogicalVariable> variables = new ArrayList<LogicalVariable>();

		// Schema mapping a String to LogicalVariable, Linked Hash Map allows to
		// keep the same order to the query.
		Map<String, LogicalVariable> fieldNameMap = new LinkedHashMap<String, LogicalVariable>();

		for (LogicalFieldSchema field : fields) {

			types.add(field.type);

			LogicalVariable var = new LogicalVariable(currentVariable++);
			variables.add(var);
			fieldNameMap.put(field.alias, var);
		}

		// TODO: Extract the path to the file to load, instead of hard-coded
		// path.
		PigDataSource dataSource = new PigDataSource(load.getSchemaFile(),
				types.toArray());
		ILogicalOperator currentOperator = new DataSourceScanOperator(
				variables, dataSource);

		// LOAD operator is the leaf of the plan and therefore has to contain an
		// emptryTupleOperator in its input
		ILogicalOperator ets = new EmptyTupleSourceOperator();
		currentOperator.getInputs().add(
				new MutableObject<ILogicalOperator>(ets));

		// currentOperator is the statement in our case it's the load
		Pair<ILogicalOperator, Map<String, LogicalVariable>> p = Pair.of(
				currentOperator, fieldNameMap);
		currentPlan.put(load.getAlias(), p);

		System.out.println("Visiting the Load node");

	}

	public void visit(LOFilter filter) throws FrontendException {

		System.out.println("Visiting the Filter node");
	}

	public void visit(LOStore store) throws FrontendException {

		String destination = store.getOutputSpec().getFileName();
		IDataSink dataSink = new PigDataStore<String>(destination);

		String alias = store.getAlias();
		Pair<ILogicalOperator, Map<String, LogicalVariable>> p = currentPlan
				.get(alias);
		List<Mutable<ILogicalExpression>> expressions = new ArrayList<Mutable<ILogicalExpression>>();

		for (LogicalVariable v : p.getRight().values()) {
			expressions.add(new MutableObject<ILogicalExpression>(
					new VariableReferenceExpression(v)));
		}

		WriteOperator write = new WriteOperator(expressions, dataSink);

		// inside the write statement, there is only the destination where to
		// store
		// and the list of variables. We still need to add the tuples that are
		// used as a context (otherwise it's meaningless)

		write.getInputs().add(new MutableObject<ILogicalOperator>(p.getLeft()));
		roots.add(new MutableObject<ILogicalOperator>(write));
		System.out.println("Visiting the Store node");
	}

	public void visit(LOJoin join) throws FrontendException {

		List<Operator> predecessors = join.getPlan().getPredecessors(join);
		// First list is the tables, second list are the fields to join by for
		// one particular table

		List<ILogicalOperator> joinOps = new ArrayList<ILogicalOperator>();
		List<List<LogicalVariable>> joinColumns = new ArrayList<List<LogicalVariable>>();
		Map<String, LogicalVariable> joinSchema = new LinkedHashMap<String, LogicalVariable>();
		int childCounter = 0;

		for (Operator op : predecessors) {
			LogicalRelationalOperator inOp = (LogicalRelationalOperator) op;
			Pair<ILogicalOperator, Map<String, LogicalVariable>> pair = currentPlan
					.get(inOp.getAlias());
			joinOps.add(pair.getLeft());
			for(Map.Entry<String, LogicalVariable> e : pair.getRight().entrySet()) {
				joinSchema.put(inOp.getAlias() + ":" + e.getKey(), e.getValue());
			}
			Collection<LogicalExpressionPlan> exprs = join
					.getJoinPlan(childCounter);
			List<LogicalVariable> vars = new ArrayList<LogicalVariable>();
			for (LogicalExpressionPlan logExp : exprs) {
				Iterator<Operator> iterator = logExp.getOperators();
				while (iterator.hasNext()) {
					ProjectExpression operator = (ProjectExpression) iterator
							.next();
					vars.add(pair.getRight().get(
							operator.getFieldSchema().alias));
				}
			}
			joinColumns.add(vars);
			++childCounter;
		}

		ILogicalOperator result = null;
		List<LogicalVariable> prevJoinVars = null;
		int nChildren = joinOps.size();
		for (int i = 0; i < nChildren; ++i) {
			ILogicalOperator iChild = joinOps.get(i);
			List<LogicalVariable> iJoinVars = joinColumns.get(i);
			if (i == 0) {
				result = iChild;
				prevJoinVars = iJoinVars;
			} else {
				result = new InnerJoinOperator(mutable(createJoinCondition(
						prevJoinVars, iJoinVars)),
						new MutableObject<ILogicalOperator>(result),
						new MutableObject<ILogicalOperator>(iChild));
			}
		}
		currentPlan.put(join.getAlias(), Pair.of(result, joinSchema));

		System.out.println("Visiting the Join node");
	}

	private ILogicalExpression createJoinCondition(List<LogicalVariable> vars1,
			List<LogicalVariable> vars2) {

		ILogicalExpression result = null;
		int n = vars1.size();
		for (int i = 0; i < n; ++i) {
			ILogicalExpression cond = new ScalarFunctionCallExpression(
					new PigFunctionInfo(AlgebricksBuiltinFunctions.EQ),
					mutable(vre(vars1.get(i))), mutable(vre(vars2.get(i))));
			if (result == null) {
				result = cond;
			} else {
				result = new ScalarFunctionCallExpression(new PigFunctionInfo(
						AlgebricksBuiltinFunctions.AND), mutable(result),
						mutable(cond));
			}
		}
		return result;
	}

	private static ILogicalExpression vre(LogicalVariable lVar) {
		return new VariableReferenceExpression(lVar);
	}

	private static <T> Mutable<T> mutable(T o) {
		return new MutableObject<T>(o);
	}

	public void visit(LOForEach foreach) throws FrontendException {
		System.out.println("Visiting the ForEach node");
	}

	public void visit(LOGenerate gen) throws FrontendException {
		System.out.println("Visiting the Gen node");
	}

	public void visit(LOInnerLoad load) throws FrontendException {
		System.out.println("Visiting the Load inner node");

	}

	public void visit(LOCogroup loCogroup) throws FrontendException {
		System.out.println("Visiting the CoGroup node");
	}

	public void visit(LOSplit loSplit) throws FrontendException {
		System.out.println("Visiting the Split node");
	}

	public void visit(LOSplitOutput loSplitOutput) throws FrontendException {
		System.out.println("Visiting the SplitOutput node");
	}

	public void visit(LOUnion loUnion) throws FrontendException {
		System.out.println("Visiting the Union node");
	}

	public void visit(LOSort loSort) throws FrontendException {
		System.out.println("Visiting the Sort node");
	}

	public void visit(LODistinct loDistinct) throws FrontendException {
		System.out.println("Visiting the Distinct node");
	}

	public void visit(LOLimit loLimit) throws FrontendException {
		System.out.println("Visiting the Limit node");
	}

	public void visit(LOCross loCross) throws FrontendException {
		System.out.println("Visiting the Cross node");
	}

	public void visit(LOStream loStream) throws FrontendException {
		System.out.println("Visiting the Stream node");
	}

	public void visit(LONative nativeMR) throws FrontendException {
		System.out.println("Visiting the Native MR node");
	}

	public ILogicalPlan getAlgebricksPlan() {

		ALogicalPlanImpl plan = new ALogicalPlanImpl(roots);
		return plan;
	}
}
